{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gated RNNs\n",
    "\n",
    "In this notebook we present an effective strategy to resolve the vanishing gradient problem in RNNs.\n",
    "\n",
    "Previously we discussed why long-term dependencies are impossible to learn by simple RNNs. This is due to the vanishing gradient problem. An effective approach to resolve this problem is to modify the architecture of the RNN. The main idea is to use controller gates inside the RNN cell to ensure the information flow from some selective hidden states. The gate controls how much of the current state (short-term memory) to pass to the next state as compared to the value of previous state (long-term memory). We discuss two effective gated RNN cells for retaining long-range dependency.\n",
    "- Gated Recurrent Unit (GRU) 2014\n",
    "- Long Short-term Memory (LSTM) 1997\n",
    "\n",
    "Although the GRU cell was proposed much later we discuss it first due to its design simplicity.\n",
    "\n",
    "\n",
    "\n",
    "## Gated Recurrent Unit (GRU)\n",
    "\n",
    "Ideally we expect a RNN to store information of a state that will be needed in future computation (by a future state). This state information should be retained for as long it is needed. Once the information (of a past state) is used by a state (future state), i.e., the past state's utility is fulfilled, then it should be reset. \n",
    "\n",
    "We also expect that there are states that don't carry pertinent information. We don't want to use information of a current state of this type to update the previous state (that may carry information from a distant past state). Thus, we want the RNN to ignore or reduce the importance of the current state. The backpropagation algorithm should be able to grow the importance of some states while lower importance of other states. \n",
    "\n",
    "These two expectations are summarized as follows.\n",
    "\n",
    "- Remembering (past state)\n",
    "- Forgetting (current state)\n",
    "\n",
    "The GRU provides mechanisms to accomplish these two goals, as sketched in the following figure. We compare the skeleton of a GRU architecture (right figure) with the vanilla RNN architecture (left). See that GRU modifies the vanilla RNN in two ways:\n",
    "- Create a channel to flow the past **memory** or hidden state $\\pmb{H}_{t-1}$ for retaining long-term memory (Update)\n",
    "- Create a channel to reduce (reset) the influence of past and put more emphasis on the current observation (Reset)\n",
    "\n",
    "The first channel \"remembers\" the past, and the second channel \"forgets\" the past.\n",
    "\n",
    "First, we will discuss how to \"forget\" the past.\n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_GRU_HighLevel.png\" width=900, height=500>\n",
    "\n",
    "\n",
    "### GRU: Control (reset) Past and Encourage Short-Term Memory\n",
    "\n",
    "Let's see how a past state $\\pmb{H}_{t-1}$ influences the current state $\\pmb{H}_{t}$ in a vanilla RNN. We use a mini-batch input with $m$ samples and $d$ features. The weight matrices has $h$ neurons. For nonlinear activation the tanh is used in this illustration.\n",
    "\n",
    "$\\pmb{H}_{t} = tanh(\\pmb{X}_t\\pmb{W}_{xh} + \\pmb{H}_{t-1} \\pmb{W}_{hh} + \\vec{b}_h)$\n",
    "\n",
    "Now say that we want to control the past state $\\pmb{H}_{t-1}$. We could connect a logic gate with it such that the gate is either open (outputs value 1) or close (outputs value 0). Let's call it a **reset gate** and denote it by $\\pmb{R}_t$. Using this gate, we can assert our control on the past and create a new current state. We will decide later whether we will use the updated current state as the final one. So, let's call it the **candidate** current state and denote it by $\\tilde{\\pmb{H}}_t$\n",
    "\n",
    "$\\tilde{\\pmb{H}}_{t} = tanh(\\pmb{X}_t\\pmb{W}_{xh} + (\\pmb{R}_t \\otimes \\pmb{H}_{t-1}) \\pmb{W}_{hh} + \\vec{b}_h)$\n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_GRU_1.png\" width=800, height=500>\n",
    "\n",
    "- But how do we design the reset gate? \n",
    "\n",
    "Since its output ranges between 0 and 1, we can design it using a fully-connected (FC) neural network with sigmoid activation (i.e., a FC MLP). Note that sigmoid squashes the output signal between the range 0 and 1. The reset gate will use the current input and the past hidden state to determine its output (whether to produce a value closer to 0 or 1).\n",
    "\n",
    "$\\pmb{R}_{t} = \\sigma(\\pmb{X}_t\\pmb{W}_{xr} + \\pmb{H}_{t-1} \\pmb{W}_{hr} + \\vec{b}_r)$\n",
    "\n",
    "As shown in the following figure, the output of the reset gate is multiplied (element-wise since it's a matrix) with the past state $\\pmb{H}_{t-1}$. If the reset gate output is 0, then the past state is completely ignored for computing the current candidate state $\\tilde{\\pmb{H}}_{t}$.\n",
    "\n",
    "Thus, now we know how to \"forget\" the past. Let's see how to \"remember\" the past.\n",
    "\n",
    "\n",
    "\n",
    "### GRU: Remember (update) Past and Encourage Long-Term Memory\n",
    "\n",
    "\n",
    "To remember the past, we need another logic gate that will determine whether to retain the past state $\\pmb{H}_{t-1}$ or the current candidate state $\\tilde{\\pmb{H}}_{t}$. In other words, this gate will **update** the current state by either retaining the past state or using only the current observation. We call this gate as the update gate and denote it with $\\pmb{Z}_t$.\n",
    "\n",
    "$\\pmb{H}_{t} = \\pmb{Z}_t \\otimes \\pmb{H}_{t-1} + (1 - \\pmb{Z}_t)\\otimes \\tilde{\\pmb{H}}_{t}$\n",
    "\n",
    "The update gate outputs values between 0 (close) and 1 (open). For example, if $\\pmb{Z}_t = 0$, the GRU uses the candidate current state (which is based on the current observation) to update its state. In other words, it forgets the past. On the other hand, if $\\pmb{Z}_t = 1$, then GRU only retains the past. \n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_GRU_2.png\" width=800, height=500>\n",
    "\n",
    "The update gate is designed using a FC network with sigmoid activation:\n",
    "\n",
    "$\\pmb{Z}_{t} = \\sigma(\\pmb{X}_t\\pmb{W}_{xz} + \\pmb{H}_{t-1} \\pmb{W}_{hz} + \\vec{b}_z)$\n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_GRU_3.png\" width=800, height=500>\n",
    "\n",
    "The above figure shows the complete GRU architecture. We see that it enables two things.\n",
    "\n",
    "- Capture short-term dependencies via the reset gate\n",
    "- Capture long-term dependencies via the update gate\n",
    "\n",
    "The GRU equations are summarized below.\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_GRU_Equations.png\" width=600, height=400>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Long Short-Term Memory (LSTM)\n",
    "\n",
    "We will approach the LSTM architecture from our understanding of the GRU architecture. Let's compare the two sets of equations from these two architectures.\n",
    "\n",
    "The first distinctive feature of the LSTM is that its state is split into two vectors: hidden-state $\\pmb{H}_t$ and memory state or **cell** state $\\pmb{C}_t$. Here the $\\pmb{C}_t$ could be considered as the long-term state, while the $\\pmb{H}_t$ is the short-term state. \n",
    "\n",
    "A key difference between LSTM and GRU is that in LSTM the memory state (similar to the hidden state in GRU) is updated using two gates (in GRU we used only one gate).\n",
    "\n",
    "LSTM processes two types of state information.\n",
    "- Past memory state $\\pmb{C}_{t-1}$ (contains long-term memory)\n",
    "- Current (candidate) memory state $\\tilde{\\pmb{C}}_{t}$ (contains short-term memory)\n",
    "\n",
    "Then, it creates the updated memory state by determining whether to **forget** the past memory state (long-term memory) and retain the candidate memory state (short-term memory).\n",
    "\n",
    "$\\pmb{C}_t = \\pmb{F}_t \\otimes \\pmb{C}_{t-1} + \\pmb{I}_t \\otimes \\tilde{\\pmb{C}}_t$\n",
    "\n",
    "Thus, to begin with, it needs at least two gates:\n",
    "- Forget gate: to control the long-term memory\n",
    "- Input gate: to control the current state information (short-term memory)\n",
    "\n",
    "It, then, uses an output gate to determine how much of the updated memory state should be passed to the next hidden state.\n",
    "\n",
    "$\\pmb{H}_t = \\pmb{O}_t \\otimes tanh({\\pmb{C}_t})$\n",
    "\n",
    "Let's discuss these three gates.\n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_GRU_LSTM_Equations.png\" width=800, height=600>\n",
    "\n",
    "### LSTM: Forget Gate - Controls Long-Term Memory\n",
    "\n",
    "The forget gate controls the flow of the long-term memory from the memory state. It is learned by a FC network with sigmoid using the past hidden state and the current input. It outputs values between 0 and 1.\n",
    "\n",
    "$\\pmb{F}_t = \\sigma(\\pmb{X}_t \\pmb{W}_{xf} + \\pmb{H}_{t-1} \\pmb{W}_{hf} + \\vec{b}_f)$\n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_LSTM_1.png\" width=800, height=500>\n",
    "\n",
    "\n",
    "### LSTM: Input Gate - Controls Short-Term Memory\n",
    "\n",
    "\n",
    "The input gate is used to determine how much of the current state information to be retained. Thus, first we compute a candidate memory state using the past hiddent state and the current input.\n",
    "\n",
    "$\\tilde{\\pmb{C}}_t = tanh(\\pmb{X}_t \\pmb{W}_{xc} + \\pmb{H}_{t-1} \\pmb{W}_{hc} + \\vec{b}_c)$\n",
    "\n",
    "It can be considered as a state that contains only short-term information as it is based on the past hidden-state (which is a short-term state) and the current input state.\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_LSTM_2.png\" width=800, height=500>\n",
    "\n",
    "This short-term state is controlled by the input gate, which is learned by a FC network with sigmoid using the past hidden state and the current input. It outputs values between 0 and 1.\n",
    "\n",
    "$\\pmb{I}_t = \\sigma(\\pmb{X}_t \\pmb{W}_{xi} + \\pmb{H}_{t-1} \\pmb{W}_{hi} + \\vec{b}_i)$\n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_LSTM_3.png\" width=800, height=500>\n",
    "\n",
    "\n",
    "### LSTM: Retain Long-Term Memory or Short-Term Memory?\n",
    "\n",
    "Both the forget gate and input gate are used to determine whether to retain long-term or short memory in the updated memory state.\n",
    "\n",
    "$\\pmb{C}_t = \\pmb{F}_t \\otimes \\pmb{C}_{t-1} + \\pmb{I}_t \\otimes \\tilde{\\pmb{C}}_t$\n",
    "\n",
    "- If $\\pmb{F}_t = 1$ and $\\pmb{I}_t = 0$: Long-term memory would be preserved dropping short-term memory.\n",
    "- If $\\pmb{F}_t = 0$ and $\\pmb{I}_t = 1$: Long-term memory would be discarded, only short-term memory is retained.\n",
    "\n",
    "Observe that if an input is important, LSTM preserves it by the input gate. For an important input, the input gate learns a non-zero value. Then, the important input is added to the long-term memory. Thus, the input gate helps to preserve an important input. The long-term memory is passed from one timestep to the next through a long sequence without losing its memory as long as the forget gate is not set to 0. Whenever the LSTM decides not to retain a long-term memory, it can \"reset\" it by using a 0 value of the forget gate.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_LSTM_4.png\" width=800, height=500>\n",
    "\n",
    "### LSTM: Updated Hidden State by the Output Gate\n",
    "\n",
    "\n",
    "Finally, the output gate controls which parts of the long-term state should be read and output at this time step, both to $\\pmb{H}_t$ and to $\\pmb{Y}_t$.\n",
    "\n",
    "$\\pmb{H}_t = \\pmb{O}_t \\otimes tanh({\\pmb{C}_t})$\n",
    "\n",
    "Just like the input and forget gates, the output gate is learned by a FC network with sigmoid using the past hidden state and the current input. It outputs values between 0 and 1.\n",
    "\n",
    "$\\pmb{O}_t = \\sigma(\\pmb{X}_t \\pmb{W}_{xo} + \\pmb{H}_{t-1} \\pmb{W}_{ho} + \\vec{b}_o)$\n",
    "\n",
    "\n",
    "Whenever the output gate is 1, the LSTM effectively passes all memory information through to the predictor, whereas for output 0, it retains all the information only within the memory cell and perform no further processing.\n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_LSTM_5.png\" width=800, height=500>\n",
    "\n",
    "\n",
    "\n",
    "In short, an LSTM cell can learn to recognize an important input (that’s the role of the input gate), store it in the long-term state, preserve it for as long as it is needed (that’s the role of the forget gate), and extract it whenever it is needed. This explains why these cells have been amazingly successful at capturing long-term patterns in time series, long texts, audio recordings, and more."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How does the Gated RNNs solve the vanishing gradient problem?\n",
    "\n",
    "\n",
    "We use the LSTM architecture to show how it prevents the vanishing gradient problem. We start by writing the expression for computing the loss gradient with respect to $\\pmb{W}$ for the vanilla RNNs, \n",
    "\n",
    "$\\frac{\\partial \\mathcal{L}}{\\pmb{W}} = \\sum_{t = 0}^{T} \\frac{\\partial l(\\vec{y_t}, \\vec{y_t}^*)}{\\partial \\vec{y_t}} \\frac{\\partial \\vec{y_t}}{\\partial \\vec{z_t}} \\frac{\\partial \\vec{z_t}}{\\partial \\vec{h}_t} \\frac{\\partial \\vec{h}_t}{\\partial \\pmb{W}}$\n",
    "\n",
    "where \n",
    "\n",
    "$ \\pmb{W} =\n",
    "\\begin{bmatrix}\n",
    "  \\pmb{W}_{xh}  \\\\\n",
    "  \\pmb{W}_{hh} \n",
    " \\end{bmatrix}\n",
    "$\n",
    "\n",
    "\n",
    "In LSTMs, past hidden states $\\vec{h}_{t-1}$ do not directly pass to the current hidden state $\\vec{h}_t$. Instead the past hidden state $\\vec{h}_{t-1}$ and the past cell state or memory state $\\pmb{c}_{t-1}$ are combined to $\\vec{c}_t$, which produces $\\vec{h}_t$. Thus, during backward propagation, $\\vec{h}_t$ passes backward through the cell state $\\pmb{c}_{t}$. Including the cell state $\\pmb{c}_{t}$ in the loss gradient equation, we get:\n",
    "\n",
    "$\\frac{\\partial \\mathcal{L}}{\\pmb{W}} = \\sum_{t = 0}^{T} \\frac{\\partial l(\\vec{y_t}, \\vec{y_t}^*)}{\\partial \\vec{y_t}} \\frac{\\partial \\vec{y_t}}{\\partial \\vec{z_t}} \\frac{\\partial \\vec{z_t}}{\\partial \\vec{h}_t} \\frac{\\partial \\vec{h}_t}{\\partial \\pmb{c}_t}\\frac{\\partial \\vec{c}_t}{\\partial \\pmb{W}}$\n",
    "\n",
    "\n",
    "\n",
    "=> $\\frac{\\partial \\mathcal{L}}{\\pmb{W}} = \\sum_{t = 0}^{T} \\frac{\\partial l(\\vec{y_t}, \\vec{y_t}^*)}{\\partial \\vec{y}_t} [\\phi^{'}\\otimes(\\vec{z}_t)] \\pmb{W}_{hy} \\frac{\\partial \\vec{h}_t}{\\partial \\pmb{c}_t}\\frac{\\partial \\vec{c}_t}{\\partial \\pmb{W}}$\n",
    "\n",
    "Let's focus on the last term $\\frac{\\partial \\vec{c}_t}{\\partial \\pmb{W}}$. The current memory cell state $\\vec{c}_t$ is influenced by the previous state $\\vec{c}_{t-1}$, which is influenced by its previous state $\\vec{c}_{t-2}$, and so on until it reaches to the cell state $\\vec{c}_{0}$ at timestep 0 forming a chain of influence. Including this chain effect, we obtain:\n",
    "\n",
    "$\\frac{\\partial \\vec{c}_t}{\\partial \\pmb{W}} = \\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}} \\frac{\\partial \\vec{c}_{t-1}}{\\partial \\vec{c}_{t-1}} ... \\frac{\\partial \\vec{c}_2}{\\partial \\vec{c}_1}\\frac{\\partial \\vec{c}_1}{\\partial \\vec{c}_0} \\frac{\\partial \\vec{c}_0}{\\partial \\pmb{W}}$\n",
    "\n",
    "\n",
    "=> $\\frac{\\partial \\vec{c}_t}{\\partial \\pmb{W}} = (\\prod_{t=1}^t \\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}}) \\frac{\\partial \\vec{c}_0}{\\partial \\pmb{W}}$\n",
    "\n",
    "\n",
    "\n",
    "This product is reminiscent of the gradient term in vanilla RNNs that led to the exploding/vanishing gradient problem . We now show why this product term will not cause the same issue in LSTMs. Let's write the equation for computing $\\vec{c}_t$.\n",
    "\n",
    "$\\vec{c}_t = \\vec{c}_{t-1} \\otimes  \\vec{f}_t + \\tilde{\\vec{c}}_t \\otimes \\vec{i}_t $\n",
    "\n",
    "Where,\n",
    "\n",
    "\n",
    "\n",
    "$\\vec{i}_t = \\sigma( \\pmb{W}_{xi}^T \\vec{x}_t + \\pmb{W}_{hi}^T \\vec{h}_{t-1}  + \\vec{b}_i)$\n",
    "\n",
    "$\\vec{f}_t = \\sigma(\\pmb{W}_{xf}^T \\vec{x}_t  +  \\pmb{W}_{hf}^T \\vec{h}_{t-1} + \\vec{b}_f)$\n",
    "\n",
    "$\\tilde{\\vec{c}}_t = tanh( \\pmb{W}_{xc}^T \\vec{x}_t +  \\pmb{W}_{hc}^T \\vec{h}_{t-1} + \\vec{b}_c)$\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "Let's compute $\\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}}$.\n",
    "\n",
    "\n",
    "\n",
    "$\\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}} = \\frac{\\partial [\\vec{c}_{t-1} \\otimes  \\vec{f}_t + \\tilde{\\vec{c}}_t \\otimes \\vec{i}_t]}{\\partial \\vec{c}_{t-1}} $\n",
    "\n",
    "=> $\\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}} = \\frac{\\partial [\\vec{c}_{t-1} \\otimes  \\vec{f}_t ]}{\\partial \\vec{c}_{t-1}} + \\frac{\\partial [ \\tilde{\\vec{c}}_t \\otimes \\vec{i}_t]}{\\partial \\vec{c}_{t-1}} $\n",
    "\n",
    "\n",
    "=> $\\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}} = \\frac{\\partial \\vec{f}_t}{\\partial \\vec{c}_{t-1}} \\vec{c}_{t-1} + \\frac{\\partial \\vec{c}_{t-1}}{\\partial \\vec{c}_{t-1}} \\vec{f}_{t} + \\frac{\\partial \\vec{i}_t}{\\partial \\vec{c}_{t-1}} \\tilde{\\vec{c}}_{t} + \\frac{\\partial \\tilde{\\vec{c}}_t}{\\partial \\vec{c}_{t-1}} \\vec{i}_{t}$\n",
    "\n",
    "\n",
    "=> $\\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}} = \\frac{\\partial \\vec{f}_t}{\\partial \\vec{c}_{t-1}} \\vec{c}_{t-1} + \\vec{f}_{t} + \\frac{\\partial \\vec{i}_t}{\\partial \\vec{c}_{t-1}} \\tilde{\\vec{c}}_{t} + \\frac{\\partial \\tilde{\\vec{c}}_t}{\\partial \\vec{c}_{t-1}} \\vec{i}_{t}$\n",
    "\n",
    "\n",
    "We compute the 1st, 3rd and 4th term as follows.\n",
    "\n",
    "- 1st term:\n",
    "$\\frac{\\partial \\vec{f}_t}{\\partial \\vec{c}_{t-1}} \\vec{c}_{t-1} = \\frac{\\partial [\\sigma(\\pmb{W}_{xf}^T \\vec{x}_t  +  \\pmb{W}_{hf}^T \\vec{h}_{t-1} + \\vec{b}_f)]}{\\partial \\vec{c}_{t-1}} \\vec{c}_{t-1}$\n",
    "\n",
    "=> $\\frac{\\partial \\vec{f}_t}{\\partial \\vec{c}_{t-1}} \\vec{c}_{t-1} = \\sigma^{'}(\\pmb{W}_{xf}^T \\vec{x}_t  +  \\pmb{W}_{hf}^T \\vec{h}_{t-1} + \\vec{b}_f) \\pmb{W}_{hf}^T \\frac{\\partial \\vec{h}_{t-1}}{\\partial \\vec{c}_{t-1}} \\vec{c}_{t-1}$\n",
    "\n",
    "\n",
    "=> $\\frac{\\partial \\vec{f}_t}{\\partial \\vec{c}_{t-1}}\\vec{c}_{t-1} = \\sigma^{'}(\\pmb{W}_{xf}^T \\vec{x}_t  +  \\pmb{W}_{hf}^T \\vec{h}_{t-1} + \\vec{b}_f) \\pmb{W}_{hf}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\vec{c}_{t-1}$\n",
    "\n",
    "=> $\\frac{\\partial \\vec{f}_t}{\\partial \\vec{c}_{t-1}} \\vec{c}_{t-1} = \\sigma^{'}(.) \\pmb{W}_{hf}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\vec{c}_{t-1}$\n",
    "\n",
    "\n",
    "The 3rd and the 4th terms can be computed in the exact same fashion.\n",
    "\n",
    "- 3rd term:\n",
    "\n",
    "$\\frac{\\partial \\vec{i}_t}{\\partial \\vec{c}_{t-1}} \\tilde{\\vec{c}}_{t} =  \\sigma^{'}(\\pmb{W}_{xi}^T \\vec{x}_t  +  \\pmb{W}_{hi}^T \\vec{h}_{t-1} + \\vec{b}_f) \\pmb{W}_{hi}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\tilde{\\vec{c}}_{t}$\n",
    "\n",
    "=> $\\frac{\\partial \\vec{i}_t}{\\partial \\vec{c}_{t-1}} \\tilde{\\vec{c}}_{t} =  \\sigma^{'}(.) \\pmb{W}_{hi}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\tilde{\\vec{c}}_{t}$\n",
    "\n",
    "\n",
    "- 4th term:\n",
    "\n",
    "$\\frac{\\partial \\tilde{\\vec{c}}_t}{\\partial \\vec{c}_{t-1}} \\vec{i}_{t} =  \\sigma^{'}(\\pmb{W}_{xc}^T \\vec{x}_t  +  \\pmb{W}_{hc}^T \\vec{h}_{t-1} + \\vec{b}_f) \\pmb{W}_{hc}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\vec{i}_{t}$\n",
    "\n",
    "=> $\\frac{\\partial \\tilde{\\vec{c}}_t}{\\partial \\vec{c}_{t-1}} \\vec{i}_{t} =  \\sigma^{'}(.) \\pmb{W}_{hc}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\vec{i}_{t}$\n",
    "\n",
    "Using the above 3 expressions we finally get:\n",
    "\n",
    "\n",
    "=> $\\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}} = \\sigma^{'}(.) \\pmb{W}_{hf}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\vec{c}_{t-1} + \\vec{f}_{t} + \\sigma^{'}(.) \\pmb{W}_{hi}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\tilde{\\vec{c}}_{t} + \\sigma^{'}(.) \\pmb{W}_{hc}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\vec{i}_{t}$\n",
    "\n",
    "\n",
    "The above expression allows us to complete the loss gradient.\n",
    "\n",
    "\n",
    "$\\frac{\\partial \\mathcal{L}}{\\pmb{W}} = \\sum_{t = 0}^{T} \\frac{\\partial l(\\vec{y_t}, \\vec{y_t}^*)}{\\partial \\vec{y}_t} [\\phi^{'}\\otimes(\\vec{z}_t)] \\pmb{W}_{hy} \\frac{\\partial \\vec{h}_t}{\\partial \\pmb{c}_t}(\\prod_{t=1}^t \\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}}) \\frac{\\partial \\vec{c}_0}{\\partial \\pmb{W}}$\n",
    "\n",
    "=> $\\frac{\\partial \\mathcal{L}}{\\pmb{W}} = \\sum_{t = 0}^{T} \\frac{\\partial l(\\vec{y_t}, \\vec{y_t}^*)}{\\partial \\vec{y}_t} [\\phi^{'}\\otimes(\\vec{z}_t)] \\pmb{W}_{hy} \\frac{\\partial \\vec{h}_t}{\\partial \\pmb{c}_t}(\\prod_{t=1}^t \\sigma^{'}(.) \\pmb{W}_{hf}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\vec{c}_{t-1} + \\vec{f}_{t} + \\sigma^{'}(.) \\pmb{W}_{hi}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\tilde{\\vec{c}}_{t} + \\sigma^{'}(.) \\pmb{W}_{hc}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\vec{i}_{t}) \\frac{\\partial \\vec{c}_0}{\\partial \\pmb{W}}$\n",
    "\n",
    "\n",
    "\n",
    "Let's see why $\\prod_{t=1}^T\\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}}$ will not lead to the vanishing gradient problem. We use a shorthand notation to represent this product.\n",
    "\n",
    "\n",
    "$\\prod_{t=1}^T \\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}} = \\prod_{t=1}^T [X + \\sigma(.) + Y + Z]$\n",
    "\n",
    "Here,\n",
    "\n",
    "$X = \\sigma^{'}(.) \\pmb{W}_{hf}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\vec{c}_{t-1}$\n",
    "\n",
    "$\\sigma(.) = \\vec{f}_{t} = \\sigma(\\pmb{W}_{xf}^T \\vec{x}_t  +  \\pmb{W}_{hf}^T \\vec{h}_{t-1} + \\vec{b}_f) $\n",
    "\n",
    "$Y = \\sigma^{'}(.) \\pmb{W}_{hi}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\tilde{\\vec{c}}_{t}$\n",
    "\n",
    "$Z = + \\sigma^{'}(.) \\pmb{W}_{hc}^T \\vec{o}_{t-1} \\otimes tanh^{'}(\\vec{c}_{t-1})\\vec{i}_{t}$\n",
    "\n",
    "\n",
    "Below we compare LSTM's $\\prod_{t=1}^T \\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}}$ term with a similar expression $\\prod_{t=1}^T \\frac{\\partial \\vec{h}_t}{\\partial \\vec{h}_{t-1}}$ of vanilla RNNs that was responsible for the driving the gradient to zero in vanilla RNNs.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_VanishingGradient.png\" width=600, height=400>\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_LSTM_VanishingGradient.png\" width=1000, height=800>\n",
    "\n",
    "Observe that in vanilla RNNs, the derivative of the activation function (assuming tanh activation) is always between 0 and 1 (as shown below). \n",
    "- If W > 1, but $tanh^{'}(.) < 1$: successive products of these two terms will lead to zero\n",
    "- If W < 1, but $tanh^{'}(.) \\leq 1$: successive products of these two terms will lead to zero\n",
    "\n",
    "However, in LSTMs, the backpropagation through a long sequence doesn't force the successive products of $\\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}}$ to zero, i.e., gradients are not bound to vanish. Instead due to its gated control mechanism, the gradients are well-behaved. There are two main reasons.\n",
    "\n",
    "- Unlike the vanilla RNNs, $\\frac{\\partial \\vec{c}_t}{\\partial \\vec{c}_{t-1}}$ is additive. The forget gate term (2nd term) is always between 0 and 1. If the X, Y and Z terms in the product start to converge to zero, LSTM cell will set the forget gate term to 1 (highest value of sigmoid, shown below). This happens when the network wants to retain long-range dependencies by letting the cell state to pass through a long sequence. In such a case, the near 1 value of the forget gate tends to shrink at a much slower rate than when compared to the derivative values of tanh in vanilla RNNs, which later on during the training processes, are likely to be saturated and thus have a value close to 0.\n",
    "\n",
    "- The LSTM cell learns to decide when to let the gradients vanish (i.e., X, Y, Z, and $\\sigma(.)$ are driven to 0) and when to preserve the gradients by setting the gate values accordingly.\n",
    "\n",
    "\n",
    "<img src=\"https://cse.unl.edu/~hasan/Pics/RNN_Sigmoid_Tanh.png\" width=600, height=400>\n",
    "\n",
    "Thus, LSTMs solve the vanishing gradient problem using a unique additive gradient structure that includes direct access to the forget gate’s activations. It enables the network to encourage desired behavior from the loss gradient using frequent gate updates on every timestep of backpropagation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Peephole Connections\n",
    "\n",
    "\n",
    "In a regular LSTM cell, the gate controllers can look only at the input $x_t$ and the previous short-term state $h_{t-1}$. It may be a good idea to give them a bit more context by letting them peek at the long-term state as well. This idea was proposed by Felix Gers and Jürgen Schmidhuber in 2000.\n",
    "\n",
    "https://www.jmlr.org/papers/volume3/gers02a/gers02a.pdf\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "They proposed an LSTM variant with extra connections called peephole connections: the previous long-term state $c_{t-1}$ is added as an input to the controllers of the forget gate and the input gate, and the current long-term state $c_{t}$ is added as input to the controller of the output gate. \n",
    "\n",
    "This often improves performance, but not always, and there is no clear pattern for which tasks are better off with or without them: we will have to try it on your task and see if it helps.\n",
    "\n",
    "\n",
    "In TensorFlow we can implement the peephole connection based LSTM using the experimental tf.keras.experimental.PeepholeLSTMCell.\n",
    "https://www.tensorflow.org/api_docs/python/tf/keras/experimental/PeepholeLSTMCell\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
